package com.liuhengjia.langchain4japply.api;

import com.liuhengjia.langchain4japply.config.assistant.Assistant;
import dev.langchain4j.data.document.Document;
import dev.langchain4j.data.document.loader.FileSystemDocumentLoader;
import dev.langchain4j.data.document.splitter.DocumentByLineSplitter;
import dev.langchain4j.data.segment.TextSegment;
import dev.langchain4j.model.embedding.EmbeddingModel;
import dev.langchain4j.store.embedding.EmbeddingStore;
import dev.langchain4j.store.embedding.EmbeddingStoreIngestor;
import lombok.RequiredArgsConstructor;
import lombok.extern.slf4j.Slf4j;
import org.springframework.web.bind.annotation.GetMapping;
import org.springframework.web.bind.annotation.RequestMapping;
import org.springframework.web.bind.annotation.RequestParam;
import org.springframework.web.bind.annotation.RestController;

import java.nio.file.Paths;
import java.util.List;

/**
 * RAG API
 */
@Slf4j
@RestController
@RequestMapping("/api/rag")
@RequiredArgsConstructor
public class RagAPI {
    private final EmbeddingStore<TextSegment> embeddingStore;
    private final Assistant ragAssistant;
    private final EmbeddingStore<TextSegment> pgVectorEmbeddingStore;
    private final EmbeddingModel embeddingModel;

    /**
     * 高级聊天
     *
     * @param message 用户信息
     * @return 模型响应
     */
    @GetMapping("/high/chat")
    public String lowChat(@RequestParam(value = "message") String message) {
        return ragAssistant.chat(message);
    }

    /**
     * 加载文档至向量数据库
     *
     * @return 操作结果
     */
    @GetMapping("/load")
    public String loadDocument() {
        try {
            List<Document> documentList = FileSystemDocumentLoader.loadDocuments(
                    Paths.get(getClass().getClassLoader().getResource("pdf").toURI())
                            .toAbsolutePath()
            );
            EmbeddingStoreIngestor.ingest(documentList, pgVectorEmbeddingStore);
        } catch (Exception e) {
            log.error(e.getMessage());
            return "error";
        }
        return "success";
    }

    /**
     * 加载文档至向量数据库
     *
     * @return 操作结果
     */
    @GetMapping("/load-splitter")
    public String loadDocumentWithSplitter() {
        try {
            List<Document> documentList = FileSystemDocumentLoader.loadDocuments(
                    Paths.get(getClass().getClassLoader().getResource("pdf").toURI())
                            .toAbsolutePath()
            );
            EmbeddingStoreIngestor.builder()
                    .embeddingStore(pgVectorEmbeddingStore)
                    .embeddingModel(embeddingModel)
                    .documentSplitter(new DocumentByLineSplitter(100, 40))
                    .build().ingest(documentList);
        } catch (Exception e) {
            log.error(e.getMessage());
            return "error";
        }
        return "success";
    }
}
